# %%
import math
import heapq
from tqdm import tqdm
import numpy as np

train_path = "../data/train.txt"
attribute_path = "../data/itemAttribute.txt"
test_path = "../data/test.txt"
result_path = "result.txt"

# %%
start_line = 0 
end_line = 0
user_item_list = {}
item_set = set()

with open(train_path, 'r') as file: 
    train_lines = file.readlines()

for index, line in enumerate(train_lines):
    if '|' in line:
        userid, num_rating_item = line.split('|')
        start_line = index + 1
        end_line = index + int(num_rating_item)
        user_item_list[userid] = {}
    elif index >= start_line and index <= end_line:
        line = line.strip('\n')
        item, rating = line.split("  ")
        user_item_list[userid][item] = int(rating)
        item_set.add(item)

print("Finished loading train data")

# %%
item_user_list = {}
for userid, item_list in user_item_list.items():
    for item in item_list:
        item_user_list.setdefault(item, dict())
        item_user_list[item][userid] = item_list[item]

item_rating_avg = {}
for item, user_list in item_user_list.items():
    avg = 0.0
    for user, rating_score in user_list.items():
        avg += rating_score
    avg = avg / len(user_list)
    item_rating_avg[item] = avg
    
item_len = len(item_user_list)
print("Finished calculating item rating avg")

# %%
user_rating_avg = {}
all_avg = 0.0
for userid, item_list in user_item_list.items():
    avg = 0.0
    for item, rating_score in item_list.items():
        avg += rating_score
        all_avg += rating_score
    avg = avg / len(item_list)
    user_rating_avg[userid] = avg
all_avg = all_avg / (len(train_lines) - len(user_item_list))

print("Finished calculating user rating avg and all rating avg")

# %%
item_norm = {}
for item, user_list in item_user_list.items():
    sum = 0.0
    for user, rating_score in user_list.items():
        sum += (rating_score - item_rating_avg[item]) ** 2
    sum = math.sqrt(sum)
    item_norm[item] = sum

print("Finished calculating user norm")

# %%
with open(attribute_path, 'r') as file: 
    attribute_lines = file.readlines()

item_attr_list = {}

for line in attribute_lines:
    line = line.strip('\n')
    item, attr1, attr2 = line.split("|")
    if attr1 == "None":
        attr1 = 0
    else:
        attr1 = int(attr1)
    if attr2 == "None":
        attr2 = 0
    else:
        attr2 = int(attr2)
    item_attr_list[item] = [attr1, attr2]

print("Finished loading item attribute")

# %%
similarity = {}
for i, (item1, user_list1) in tqdm(enumerate(user_item_list.items()), total=item_len):
    similarity[item1] = {}
    for j, (item2, user_list2) in enumerate(user_item_list.items()):
        if item1 not in item_user_list or item2 not in item_user_list:
            similarity[item1][item2] = 0.0
            continue
        if item_norm[item1] * item_norm[item2] == 0:
            similarity[item1][item2] = 0.0
        else:
            if i > j:
                similarity[item1][item2] = similarity[item2][item1]
            elif i == j:
                similarity[item1][item2] = 1.0
            else:
                cos_sim = 0.0
                for user1, rating in user_list1.items():
                    if user1 in user_list2:
                        cos_sim += (rating - item_rating_avg[item1]) * (user_list2[user1] - user_rating_avg[item2])
                cos_sim = cos_sim / (item_norm[item1] * item_norm[item2])
                similarity[item1][item2] = cos_sim
        attr_sim = 0.0
        if item1 in item_attr_list and item2 in item_attr_list:
            attr_sim = (item_attr_list[item1][0] * item_attr_list[item2][0]) + (item_attr_list[item1][1] * item_attr_list[item2][1])
            attr_len = math.sqrt(item_attr_list[item1][0] ** 2 + item_attr_list[item1][1] ** 2)*math.sqrt(item_attr_list[item2][0] **2 + item_attr_list[item2][1] ** 2)
            if attr_len == 0:
                attr_sim = 0.0
            else:
                attr_sim = attr_sim / attr_len
                similarity[item1][item2] = (similarity[item1][item2] + attr_sim) / 2

# %%
with open(test_path, 'r') as file: 
    test_lines = file.readlines()

start_line = 0 
end_line = 0
test_user_item_list = {}
    
for index, line in enumerate(test_lines):
    if '|' in line:
        userid, num_rating_item = line.split('|')
        start_line = index + 1
        end_line = index + int(num_rating_item)
        test_user_item_list[userid] = {}
    elif index >= start_line and index <= end_line:
        line = line.strip('\n')
        test_user_item_list[userid][line] = int(-1)

# %%
for userid, item_list in tqdm(test_user_item_list.items(), total=len(test_user_item_list)):
    for item in item_list:
        largest_n = heapq.nlargest(5, ((k, v) for k, v in similarity[item].items() if k in user_item_list[item] and k != item and similarity[item][k] > 0.0), key=lambda item: item[1])
        rating_sum = 0.0
        sim_sum = 0.0
        for item_sim, value in largest_n:
            rating_sum = rating_sum + similarity[item][item_sim] * user_item_list[userid][item_sim]
            sim_sum = sim_sum + similarity[item][item_sim]
        if sim_sum == 0:
            if item in item_set:
                test_user_item_list[userid][item] = user_rating_avg[userid] - all_avg + item_rating_avg[item]
            else:
                test_user_item_list[userid][item] = item_rating_avg[userid]
        else:
            test_user_item_list[userid][item] = (rating_sum / sim_sum) 
        if test_user_item_list[userid][item] < 0:
            test_user_item_list[userid][item] = user_rating_avg[userid] - all_avg + item_rating_avg[item]
        if test_user_item_list[userid][item] > 100:
            test_user_item_list[userid][item] = 100


